import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt

# 生成一些示例数据：两个高斯分布的数据
np.random.seed(0)
data = np.concatenate([
    np.random.normal(loc=0.0, scale=1.0, size=200),
    np.random.normal(loc=5.0, scale=1.0, size=200)
])

# 初始化参数
mu1 = np.random.normal()
mu2 = np.random.normal()
sigma1 = np.abs(np.random.normal())
sigma2 = np.abs(np.random.normal())
pi = 0.5  # 混合系数

# EM算法
max_iter = 100
epsilon = 1e-4  # 收敛阈值

for i in range(max_iter):
    # E-step: 计算后验概率
    likelihood1 = norm.pdf(data, mu1, sigma1)
    likelihood2 = norm.pdf(data, mu2, sigma2)

    weight_likelihood = pi * likelihood1 + (1 - pi) * likelihood2
    gamma1 = (pi * likelihood1) / weight_likelihood
    gamma2 = ((1 - pi) * likelihood2) / weight_likelihood

    # M-step: 更新参数
    mu1_new = np.sum(gamma1 * data) / np.sum(gamma1)
    mu2_new = np.sum(gamma2 * data) / np.sum(gamma2)

    sigma1_new = np.sqrt(np.sum(gamma1 * (data - mu1_new) ** 2) / np.sum(gamma1))
    sigma2_new = np.sqrt(np.sum(gamma2 * (data - mu2_new) ** 2) / np.sum(gamma2))

    pi_new = np.mean(gamma1)

    # 可视化当前迭代的拟合结果
    if i % 5 == 0 or i == max_iter - 1:
        x = np.linspace(min(data), max(data), 1000)
        y_total = pi * norm.pdf(x, mu1, sigma1) + (1 - pi) * norm.pdf(x, mu2, sigma2)

        plt.hist(data, bins=30, density=True, alpha=0.6, color='gray', label="Data")
        plt.plot(x, pi * norm.pdf(x, mu1, sigma1), 'r--', label=f"Gaussian 1 (μ={mu1:.2f}, σ={sigma1:.2f})")
        plt.plot(x, (1 - pi) * norm.pdf(x, mu2, sigma2), 'g--', label=f"Gaussian 2 (μ={mu2:.2f}, σ={sigma2:.2f})")
        plt.plot(x, y_total, 'b-', label="Fitted GMM")
        plt.title(f"EM Iteration {i}")
        plt.legend()
        plt.show()

    # 检查是否收敛
    if np.abs(mu1_new - mu1) < epsilon and \
            np.abs(mu2_new - mu2) < epsilon and \
            np.abs(sigma1_new - sigma1) < epsilon and \
            np.abs(sigma2_new - sigma2) < epsilon and \
            np.abs(pi_new - pi) < epsilon:
        print(f"Converged after {i + 1} iterations")
        break

    # 更新参数
    mu1, mu2, sigma1, sigma2, pi = mu1_new, mu2_new, sigma1_new, sigma2_new, pi_new

print("Estimated parameters:")
print(f"mu1: {mu1:.2f}, sigma1: {sigma1:.2f}")
print(f"mu2: {mu2:.2f}, sigma2: {sigma2:.2f}")
print(f"pi: {pi:.2f}")
